from langchain.chains import LLMChain
from langchain.chains.sequential import SequentialChain
from langchain_core.prompts import PromptTemplate
from langchain_openai import OpenAI

#
llm = OpenAI(
    api_key='sk-VowKQBUMIkSND8WScNJtDLqf3FyqWHQ43LMVUXH1m6GZaopA',
    base_url='https://ai.nengyongai.cn/v1',
    temperature=0
)


def test():
    # 第一个 chain: 生成花的介绍
    introduction_chain = LLMChain(
        llm=llm,
        prompt=PromptTemplate.from_template('生成花的介绍'),
        output_key="introduction"
    )
    # 第二个 chain: 根据介绍写评论
    review_chain = LLMChain(
        llm=llm,
        prompt=PromptTemplate.from_template('根据介绍写评论'),
        output_key="review"
    )
    # 第三个 chain: 根据介绍和评论写推文
    social_chain = LLMChain(
        llm=llm,
        prompt=PromptTemplate.from_template('根据介绍和评论写推文'),
        output_key="social_post"
    )

    # 将三个 chain 串联起来
    all_chain = SequentialChain(
        chains=[introduction_chain, review_chain],
        input_variables=["flower_name"],
        output_variables=["introduction", "review"]
    )

    # 调用 chain
    result = all_chain('玫瑰花')
    print(result)
